# -*- coding: utf-8 -*-
"""
@Time 2020/12/14 14:18
@Author 鹄望潇湘
@File FCN_VGG16_Trainer.py
@Desc 此文件中包含训练FCN_VGG16神经网络的训练器，训练细节和逻辑控制均在此文件中
"""

import torch
from tools_zsj.Trainer import Trainer, TrainParameter
from tools_zsj.ZsjDataSet import MultiDataSet
from FCN_VGG16 import FCN_VGG16
from torch.utils.data import DataLoader
from FCN_VGG16_Tester import PixelEvaluator
import numpy as np


class FCN_VGG16_Trainer(Trainer):
    """
    这是一个为FCN_VGG16网络编写的训练器，大部分的网络训练逻辑控制均在此类中。

    :param train_datasets : 训练所需的训练数据集
    :param fcn_net: 一个FCN_VGG16的网络实例对象
    :param validate_datasets: 验证结果所需的数据集
    :param parameters: 训练所用的一些参数，包括batch_size, epoch， device等等

    """
    def __init__(self, fcn_net: FCN_VGG16, train_datasets: MultiDataSet, validate_datasets: MultiDataSet,
                 parameters: TrainParameter):
        super(FCN_VGG16_Trainer, self).__init__(fcn_net, train_datasets, parameters)
        self.validate_dataset = validate_datasets
        self.module.to(self.train_parameter.device)
        self.best_accuracy = 0.0

    def train_detail(self, index: int, item: list):
        image_data = item[0]
        correct_label = item[1]

        self.module.optimizer.zero_grad()
        predict = self.module(image_data.to(self.train_parameter.device))
        loss = self.module.loss_function(predict, correct_label.to(self.train_parameter.device))

        loss.backward()

        self.module.optimizer.step()

        if (index + 1) % 10 == 0:
            if loss.item() > 2.0:
                self.module.optimizer.__setattr__("lr", 0.1)
            elif 2.0 > loss.item() > 1.0:
                self.module.optimizer.__setattr__("lr", 0.01)
            elif 1.0 > loss.item() > 0.15:
                self.module.optimizer.__setattr__("lr", 0.0001)
            else:
                self.module.optimizer.__setattr__("lr", 0.00001)

        if (index+1) % 100 == 0:
            self.logger.info("batch: %d, loss: %.3f" % (index, loss))

    def validate_and_save(self, epoch: int):
        """
        模型参数的评估与模型参数的保存函数。从验证数据集中取出10张图片逐一进行语义分割，将分割结果与
        正确结果做对比，求出预测准确的像素点总数，从而计算准确率

        :param batch: 当前训练的batch数
        :return: 无
        """
        self.logger.info("Start evaluate...")
        self.module.eval()

        evaluator = PixelEvaluator(22)
        with torch.no_grad():
            validate_loader = DataLoader(self.validate_dataset, batch_size=5, num_workers=5, shuffle=True)
            for index, item in enumerate(validate_loader):
                image_data = item[0]
                correct_label = item[1]

                predict = self.module(image_data.to(self.train_parameter.device))
                predict = torch.max(predict, dim=1)[1]
                evaluator.push_batch(predict, correct_label)

        self.module.train()

        this_acc = evaluator.get_mean_IU()
        self.logger.info("pixel accuracy: %.4f mean accuracy: %.4f , mean IU: %.4f" %
                         (evaluator.get_correct_accuracy(), evaluator.get_mean_accuracy(), this_acc))
        if this_acc > self.best_accuracy:
            self.best_accuracy = this_acc
            torch.save(self.module.state_dict(), "./checkpoint7/checkpoint_best_7.pth")
            self.logger.info("模型参数已保存:%.4f" % this_acc)

    def epoch_finish(self, epoch: int):
        self.logger.info("第%d个epoch已完成" % epoch)
        self.validate_and_save(epoch)
        if (epoch+1) % 5 == 0:
            torch.save(self.module.state_dict(), "./checkpoint7/checkpoint_7_"+str(epoch)+".pth")
            self.logger.info("checkpoint has saved")


